// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/plan/presto_sql.h>

#include <pollux/plan/reference_query_runner.h>
#include <pollux/functions/prestosql/types/json_type.h>

namespace kumo::pollux::plan {
    namespace {
        std::string joinKeysToSql(
            const std::vector<core::FieldAccessTypedExprPtr> &keys) {
            std::vector<std::string> keyNames;
            keyNames.reserve(keys.size());
            for (const core::FieldAccessTypedExprPtr &key: keys) {
                keyNames.push_back(key->name());
            }
            return melon::join(", ", keyNames);
        }

        std::string filterToSql(const core::TypedExprPtr &filter) {
            auto call = std::dynamic_pointer_cast<const core::CallTypedExpr>(filter);
            return toCallSql(call);
        }

        std::string joinConditionAsSql(const core::AbstractJoinNode &joinNode) {
            std::stringstream out;
            for (auto i = 0; i < joinNode.leftKeys().size(); ++i) {
                if (i > 0) {
                    out << " AND ";
                }
                out << joinNode.leftKeys()[i]->name() << " = "
                        << joinNode.rightKeys()[i]->name();
            }
            if (joinNode.filter()) {
                if (!joinNode.leftKeys().empty()) {
                    out << " AND ";
                }
                out << filterToSql(joinNode.filter());
            }
            return out.str();
        }

        std::optional<std::string> toJoinSourceSql(std::optional<std::string> &&sql) {
            return sql.has_value() && sql->find(" ") != std::string::npos
                       ? fmt::format("({})", *sql)
                       : sql;
        }
    } // namespace

    void appendComma(int32_t i, std::stringstream &sql) {
        if (i > 0) {
            sql << ", ";
        }
    }

    // Returns the SQL string of the given type.
    std::string toTypeSql(const TypePtr &type) {
        switch (type->kind()) {
            case TypeKind::ARRAY:
                return fmt::format("ARRAY({})", toTypeSql(type->childAt(0)));
            case TypeKind::MAP:
                return fmt::format(
                    "MAP({}, {})",
                    toTypeSql(type->childAt(0)),
                    toTypeSql(type->childAt(1)));
            case TypeKind::ROW: {
                const auto &rowType = type->as_row();
                std::stringstream sql;
                sql << "ROW(";
                for (auto i = 0; i < type->size(); ++i) {
                    appendComma(i, sql);
                    // TODO Field names may need to be quoted.
                    sql << rowType.nameOf(i) << " " << toTypeSql(type->childAt(i));
                }
                sql << ")";
                return sql.str();
            }
            default:
                if (type->isPrimitiveType()) {
                    return type->name();
                }
                POLLUX_UNSUPPORTED("Type is not supported: {}", type->toString());
        }
    }

    std::string toLambdaSql(const core::LambdaTypedExprPtr &lambda) {
        std::stringstream sql;
        const auto &signature = lambda->signature();

        sql << "(";
        for (auto j = 0; j < signature->size(); ++j) {
            appendComma(j, sql);
            sql << signature->nameOf(j);
        }

        sql << ") -> ";
        toCallInputsSql({lambda->body()}, sql);
        return sql.str();
    }

    namespace {
        std::string toDereferenceSql(const core::DereferenceTypedExpr &dereference) {
            std::stringstream sql;
            toCallInputsSql(dereference.inputs(), sql);
            sql << "." << dereference.name();
            return sql.str();
        }
    } // namespace

    void toCallInputsSql(
        const std::vector<core::TypedExprPtr> &inputs,
        std::stringstream &sql) {
        for (auto i = 0; i < inputs.size(); ++i) {
            appendComma(i, sql);

            const auto &input = inputs.at(i);
            if (auto field =
                    std::dynamic_pointer_cast<const core::FieldAccessTypedExpr>(
                        input)) {
                sql << field->name();
            } else if (
                auto call =
                        std::dynamic_pointer_cast<const core::CallTypedExpr>(input)) {
                sql << toCallSql(call);
            } else if (
                auto lambda =
                        std::dynamic_pointer_cast<const core::LambdaTypedExpr>(input)) {
                sql << toLambdaSql(lambda);
            } else if (
                auto constantArg =
                        std::dynamic_pointer_cast<const core::ConstantTypedExpr>(input)) {
                sql << toConstantSql(*constantArg);
            } else if (
                auto castArg =
                        std::dynamic_pointer_cast<const core::CastTypedExpr>(input)) {
                sql << toCastSql(*castArg);
            } else if (
                auto concatArg =
                        std::dynamic_pointer_cast<const core::ConcatTypedExpr>(input)) {
                sql << toConcatSql(*concatArg);
            } else if (
                auto dereferenceArg =
                        std::dynamic_pointer_cast<const core::DereferenceTypedExpr>(
                            input)) {
                sql << toDereferenceSql(*dereferenceArg);
            } else {
                POLLUX_NYI("Unsupported input expression: {}.", input->toString());
            }
        }
    }

    // Returns a mapping from Pollux function names to the corresponding unary
    // operators supported in Presto SQL.
    const melon::F14FastMap<std::string, std::string> &unaryOperatorMap() {
        static melon::F14FastMap<std::string, std::string> unaryOperatorMap{
            {"negate", "-"},
            {"not", "not"},
        };
        return unaryOperatorMap;
    }

    // Returns a mapping from Pollux function names to the corresponding binary
    // operators supported in Presto SQL.
    const melon::F14FastMap<std::string, std::string> &binaryOperatorMap() {
        static melon::F14FastMap<std::string, std::string> binaryOperatorMap{
            {"plus", "+"},
            {"subtract", "-"},
            {"minus", "-"},
            {"multiply", "*"},
            {"divide", "/"},
            {"eq", "="},
            {"neq", "<>"},
            {"lt", "<"},
            {"gt", ">"},
            {"lte", "<="},
            {"gte", ">="},
            {"distinct_from", "is distinct from"},
        };
        return binaryOperatorMap;
    }

    std::string toCallSql(const core::CallTypedExprPtr &call) {
        std::stringstream sql;
        // Some functions require special SQL syntax, so handle them first.
        const auto &unaryOperators = unaryOperatorMap();
        const auto &binaryOperators = binaryOperatorMap();
        if (unaryOperators.count(call->name()) > 0) {
            POLLUX_CHECK_EQ(call->inputs().size(), 1);
            sql << "(";
            sql << fmt::format("{} ", unaryOperators.at(call->name()));
            toCallInputsSql({call->inputs()[0]}, sql);
            sql << ")";
        } else if (binaryOperators.count(call->name()) > 0) {
            POLLUX_CHECK_EQ(call->inputs().size(), 2);
            sql << "(";
            toCallInputsSql({call->inputs()[0]}, sql);
            sql << fmt::format(" {} ", binaryOperators.at(call->name()));
            toCallInputsSql({call->inputs()[1]}, sql);
            sql << ")";
        } else if (call->name() == "is_null" || call->name() == "not_null") {
            sql << "(";
            POLLUX_CHECK_EQ(call->inputs().size(), 1);
            toCallInputsSql({call->inputs()[0]}, sql);
            sql << fmt::format(" is{} null", call->name() == "not_null" ? " not" : "");
            sql << ")";
        } else if (call->name() == "in") {
            POLLUX_CHECK_GE(call->inputs().size(), 2);
            toCallInputsSql({call->inputs()[0]}, sql);
            sql << " in (";
            for (auto i = 1; i < call->inputs().size(); ++i) {
                appendComma(i - 1, sql);
                toCallInputsSql({call->inputs()[i]}, sql);
            }
            sql << ")";
        } else if (call->name() == "like") {
            sql << "(";
            toCallInputsSql({call->inputs()[0]}, sql);
            sql << " like ";
            toCallInputsSql({call->inputs()[1]}, sql);
            if (call->inputs().size() == 3) {
                sql << " escape ";
                toCallInputsSql({call->inputs()[2]}, sql);
            }
            sql << ")";
        } else if (call->name() == "or" || call->name() == "and") {
            sql << "(";
            const auto &inputs = call->inputs();
            for (auto i = 0; i < inputs.size(); ++i) {
                if (i > 0) {
                    sql << fmt::format(" {} ", call->name());
                }
                toCallInputsSql({inputs[i]}, sql);
            }
            sql << ")";
        } else if (call->name() == "array_constructor") {
            sql << "ARRAY[";
            toCallInputsSql(call->inputs(), sql);
            sql << "]";
        } else if (call->name() == "between") {
            const auto &inputs = call->inputs();
            POLLUX_CHECK_EQ(inputs.size(), 3);
            toCallInputsSql({inputs[0]}, sql);
            sql << " between ";
            toCallInputsSql({inputs[1]}, sql);
            sql << " and ";
            toCallInputsSql({inputs[2]}, sql);
        } else if (call->name() == "row_constructor") {
            sql << "row(";
            toCallInputsSql(call->inputs(), sql);
            sql << ")";
        } else if (call->name() == "subscript") {
            POLLUX_CHECK_EQ(call->inputs().size(), 2);
            toCallInputsSql({call->inputs()[0]}, sql);
            sql << "[";
            toCallInputsSql({call->inputs()[1]}, sql);
            sql << "]";
        } else {
            // Regular function call syntax.
            sql << call->name() << "(";
            toCallInputsSql(call->inputs(), sql);
            sql << ")";
        }
        return sql.str();
    }

    std::string toCastSql(const core::CastTypedExpr &cast) {
        std::stringstream sql;
        if (cast.nullOnFailure()) {
            sql << "try_cast(";
        } else {
            sql << "cast(";
        }
        toCallInputsSql(cast.inputs(), sql);
        sql << " as " << toTypeSql(cast.type());
        sql << ")";
        return sql.str();
    }

    std::string toConcatSql(const core::ConcatTypedExpr &concat) {
        std::stringstream sql;
        sql << "row(";
        toCallInputsSql(concat.inputs(), sql);
        sql << ")";
        return sql.str();
    }

    template<typename T>
    T getConstantValue(const core::ConstantTypedExpr &expr) {
        if (expr.hasValueVector()) {
            return expr.value_vector()->as<SimpleVector<T> >()->value_at(0);
        } else {
            return expr.value().value<T>();
        }
    }

    template<>
    std::string getConstantValue(const core::ConstantTypedExpr &expr) {
        if (expr.hasValueVector()) {
            return expr.value_vector()
                    ->as<SimpleVector<StringView> >()
                    ->value_at(0)
                    .getString();
        } else {
            return expr.value().value<std::string>();
        }
    }

    std::string toConstantSql(const core::ConstantTypedExpr &constant) {
        const auto &type = constant.type();
        const auto typeSql = toTypeSql(type);

        std::stringstream sql;
        if (constant.isNull()) {
            // Syntax like BIGINT 'null' for typed null is not supported, so use cast
            // instead.
            sql << fmt::format("cast(null as {})", typeSql);
        } else if (type->is_varchar() || type->is_varbinary()) {
            // Escape single quote in string literals used in SQL texts.
            if (type->is_varbinary()) {
                sql << typeSql << " ";
            }
            sql << std::quoted(getConstantValue<std::string>(constant), '\'', '\'');
        } else if (type->is_bigint()) {
            sql << getConstantValue<int64_t>(constant);
        } else if (type->isPrimitiveType()) {
            sql << fmt::format("{} '{}'", typeSql, constant.toString());
        } else {
            POLLUX_NYI(
                "Constant expressions of {} are not supported yet.", type->toString());
        }
        return sql.str();
    }

    std::string toAggregateCallSql(
        const core::CallTypedExprPtr &call,
        const std::vector<core::FieldAccessTypedExprPtr> &sortingKeys,
        const std::vector<core::SortOrder> &sortingOrders,
        bool distinct) {
        POLLUX_CHECK_EQ(sortingKeys.size(), sortingOrders.size());
        std::stringstream sql;
        sql << call->name() << "(";

        if (distinct) {
            sql << "distinct ";
        }

        toCallInputsSql(call->inputs(), sql);

        if (!sortingKeys.empty()) {
            sql << " ORDER BY ";

            for (int i = 0; i < sortingKeys.size(); i++) {
                appendComma(i, sql);
                sql << sortingKeys[i]->name() << " " << sortingOrders[i].toString();
            }
        }

        sql << ")";
        return sql.str();
    }

    void PrestoSqlPlanNodeVisitor::visit(
        const core::ValuesNode &node,
        core::PlanNodeVisitorContext &ctx) const {
        PrestoSqlPlanNodeVisitorContext &visitorContext =
                static_cast<PrestoSqlPlanNodeVisitorContext &>(ctx);
        if (!queryRunner_->isSupportedDwrfType(node.outputType())) {
            visitorContext.sql = std::nullopt;
        } else {
            visitorContext.sql = queryRunner_->getTableName(node);
        }
    }

    void PrestoSqlPlanNodeVisitor::visit(
        const core::TableScanNode &node,
        core::PlanNodeVisitorContext &ctx) const {
        static_cast<PrestoSqlPlanNodeVisitorContext &>(ctx).sql =
                node.tableHandle()->name();
    }

    void PrestoSqlPlanNodeVisitor::visit(
        const core::HashJoinNode &node,
        core::PlanNodeVisitorContext &ctx) const {
        PrestoSqlPlanNodeVisitorContext &visitorContext =
                static_cast<PrestoSqlPlanNodeVisitorContext &>(ctx);

        if (!queryRunner_->isSupportedDwrfType(node.sources()[0]->outputType()) ||
            !queryRunner_->isSupportedDwrfType(node.sources()[1]->outputType())) {
            visitorContext.sql = std::nullopt;
            return;
        }

        std::optional<std::string> probeTableName =
                toJoinSourceSql(toSql(node.sources()[0]));
        std::optional<std::string> buildTableName =
                toJoinSourceSql(toSql(node.sources()[1]));
        if (!probeTableName || !buildTableName) {
            visitorContext.sql = std::nullopt;
            return;
        }

        const auto &outputNames = node.outputType()->names();

        std::stringstream sql;
        if (node.isLeftSemiProjectJoin()) {
            sql << "SELECT "
                    << melon::join(", ", outputNames.begin(), --outputNames.end());
        } else {
            sql << "SELECT " << melon::join(", ", outputNames);
        }

        switch (node.joinType()) {
            case core::JoinType::kInner:
                sql << " FROM " << *probeTableName << " INNER JOIN " << *buildTableName
                        << " ON " << joinConditionAsSql(node);
                break;
            case core::JoinType::kLeft:
                sql << " FROM " << *probeTableName << " LEFT JOIN " << *buildTableName
                        << " ON " << joinConditionAsSql(node);
                break;
            case core::JoinType::kFull:
                sql << " FROM " << *probeTableName << " FULL OUTER JOIN "
                        << *buildTableName << " ON " << joinConditionAsSql(node);
                break;
            case core::JoinType::kLeftSemiFilter:
                // Multiple columns returned by a scalar subquery is not supported. A
                // scalar subquery expression is a subquery that returns one result row
                // from exactly one column for every input row.
                if (node.leftKeys().size() > 1) {
                    visitorContext.sql = std::nullopt;
                    return;
                }
                sql << " FROM " << *probeTableName << " WHERE "
                        << joinKeysToSql(node.leftKeys()) << " IN (SELECT "
                        << joinKeysToSql(node.rightKeys()) << " FROM " << *buildTableName;
                if (node.filter()) {
                    sql << " WHERE " << filterToSql(node.filter());
                }
                sql << ")";
                break;
            case core::JoinType::kLeftSemiProject:
                if (node.isNullAware()) {
                    sql << ", " << joinKeysToSql(node.leftKeys()) << " IN (SELECT "
                            << joinKeysToSql(node.rightKeys()) << " FROM " << *buildTableName;
                    if (node.filter()) {
                        sql << " WHERE " << filterToSql(node.filter());
                    }
                    sql << ") FROM " << *probeTableName;
                } else {
                    sql << ", EXISTS (SELECT * FROM " << *buildTableName << " WHERE "
                            << joinConditionAsSql(node);
                    sql << ") FROM " << *probeTableName;
                }
                break;
            case core::JoinType::kAnti:
                if (node.isNullAware()) {
                    sql << " FROM " << *probeTableName << " WHERE "
                            << joinKeysToSql(node.leftKeys()) << " NOT IN (SELECT "
                            << joinKeysToSql(node.rightKeys()) << " FROM " << *buildTableName;
                    if (node.filter()) {
                        sql << " WHERE " << filterToSql(node.filter());
                    }
                    sql << ")";
                } else {
                    sql << " FROM " << *probeTableName
                            << " WHERE NOT EXISTS (SELECT * FROM " << *buildTableName
                            << " WHERE " << joinConditionAsSql(node);
                    sql << ")";
                }
                break;
            default:
                POLLUX_UNREACHABLE(
                    "Unknown join type: {}", static_cast<int>(node.joinType()));
        }
        visitorContext.sql = sql.str();
    }

    void PrestoSqlPlanNodeVisitor::visit(
        const core::NestedLoopJoinNode &node,
        core::PlanNodeVisitorContext &ctx) const {
        PrestoSqlPlanNodeVisitorContext &visitorContext =
                static_cast<PrestoSqlPlanNodeVisitorContext &>(ctx);

        std::optional<std::string> probeTableName =
                toJoinSourceSql(toSql(node.sources()[0]));
        std::optional<std::string> buildTableName =
                toJoinSourceSql(toSql(node.sources()[1]));
        if (!probeTableName || !buildTableName) {
            visitorContext.sql = std::nullopt;
            return;
        }

        std::stringstream sql;
        sql << "SELECT " << melon::join(", ", node.outputType()->names());

        // Nested loop join without filter.
        POLLUX_CHECK_NULL(
            node.joinCondition(),
            "This code path should be called only for nested loop join without filter");
        const std::string joinCondition{"(1 = 1)"};
        switch (node.joinType()) {
            case core::JoinType::kInner:
                sql << " FROM " << *probeTableName << " INNER JOIN " << *buildTableName
                        << " ON " << joinCondition;
                break;
            case core::JoinType::kLeft:
                sql << " FROM " << *probeTableName << " LEFT JOIN " << *buildTableName
                        << " ON " << joinCondition;
                break;
            case core::JoinType::kFull:
                sql << " FROM " << *probeTableName << " FULL OUTER JOIN "
                        << *buildTableName << " ON " << joinCondition;
                break;
            default:
                POLLUX_UNREACHABLE(
                    "Unknown join type: {}", static_cast<int>(node.joinType()));
        }
        visitorContext.sql = sql.str();
    }

    std::optional<std::string> PrestoSqlPlanNodeVisitor::toSql(
        const core::PlanNodePtr &node) const {
        PrestoSqlPlanNodeVisitorContext sourceContext;
        node->accept(*this, sourceContext);
        return sourceContext.sql;
    }
} // namespace kumo::pollux::exec::test
